"""
处理矩阵
拿到轨迹点与各项参数，处理MPC所需的矩阵
"""
import math

import numpy as np

from bins.tool import normalize_angle
from bins.vehicleParams import VehicleParams
import scipy.sparse as sparse
import logging
import osqp
logger = logging.getLogger(__name__)


class ProcessedMatrix:
    def __init__(self):
        self.shape_u = None
        self.predicted_control = None
        self.predicted_state = None
        self.g = None
        self.H = None
        self.vehicle = None
        self.ad = None
        self.Ad = None

        self.N = None
        self.Bd = None
        self.xk = None
        self.uk = None
        self.shape = None
        self.x_ref = None
        self.Qd = None
        self.Rd = None
        self.min_u = None
        self.max_u = None
        self.v_ref = None
        self.dt = 0.1
        self.bd = np.array([[0], [0], [1]])
        self.qp = osqp.OSQP()
        self.qp_failed_time = 0
        self.qp_setup_done = False

    def init_process(self, xk: np.array = None, uk: np.array = None,  N: int = 25,
                     vehicle_conf: VehicleParams = None, v_ref: list = None, dt: float = 0.1):
        if vehicle_conf is None:
            if self.vehicle is None:
                self.vehicle = VehicleParams()
                self.vehicle.load_vehicle_params()
                self.vehicle.calculate_axis()
        else:
            self.vehicle = vehicle_conf
        if xk is not None:
            self.xk = xk
            self.shape = xk.shape[0]
        if uk is not None:
            self.uk = uk
            self.shape_u = uk.shape[1]
        if v_ref is not None:
            self.v_ref = v_ref
        self.N = N
        self.dt = dt

    def calculate_a(self, v: float, dt: float):
        four_k = v + 1 / 2 * self.dt * v ** 2 + 1 / 6 * self.dt ** 2 * v ** 3 + 1 / 24 * self.dt ** 3 * v ** 4
        self.ad = np.array([[0, four_k, 0],
                            [0, 0, four_k / self.vehicle.vehicle_params["front_axis_to_rear_axis"]],
                            [0, 0, 0]])
        tmp_a = np.eye(3) * 1 + self.ad * dt / 2
        tmp_b = np.eye(3) * 1 - self.ad * dt / 2
        self.ad = np.dot(np.linalg.inv(tmp_b), tmp_a)

    def calculate_a_lon(self):
        self.ad = np.array([[1, self.dt, 1 / 2 * self.dt ** 2],
                            [0, 1, self.dt],
                            [0, 0, 0]])
        self.bd = np.array([[0], [0], [1]])
        # tmp_a = np.eye(3) * 1 + self.ad * self.dt / 2
        # tmp_b = np.eye(3) * 1 - self.ad * self.dt / 2
        # self.ad = np.dot(np.linalg.inv(tmp_b), tmp_a)

    def calaulate_b(self, dt: float):
        self.bd = np.array([[0], [0], [1]]) * dt

    def calculate_A(self, list_v: list = None) -> bool:
        dt = 0.1
        self.calaulate_b(dt)
        self.shape = self.xk.shape[0]
        shape_u = self.uk.shape[1]
        self.shape_u = shape_u
        self.Ad = np.zeros((self.shape * self.N, self.shape))
        self.Bd = np.zeros((self.shape * self.N, self.N * shape_u))
        for i in range(self.N):
            for j in range(self.N):
                if i == 0:
                    node_j = self.bd
                else:
                    node_j = np.dot(node, self.bd)
                if i + j < self.N:
                    index = self.shape * (i + j)
                    self.Bd[index: index + self.shape, shape_u * j:shape_u * j + shape_u] = node_j
                else:
                    break
            self.calculate_a(list_v[i], dt)
            if i == 0:
                node = self.ad
            else:
                node = np.dot(self.ad, node)
            self.Ad[self.shape * i:self.shape * i + self.shape, :] = node
        qd = np.diag([80, 50, 0])
        self.Qd = np.kron(np.eye(self.N), qd)
        self.Rd = np.eye(self.N * shape_u) * 40
        return True

    def calculate_qn(self, cost_w: list):
        self.qd = np.diag(cost_w)

    def calculate_rn(self, cost_u: list):
        self.rd = np.diag(cost_u)

    def calculate_H(self):
        self.Qd = np.kron(np.eye(self.N + 1), self.qd)
        self.Rd = np.kron(np.eye(self.N), self.rd)
        self.H = np.block([[self.Qd, np.zeros((self.Qd.shape[0], self.Rd.shape[1]))],
                          [np.zeros((self.Rd.shape[1], self.Qd.shape[0])), self.Rd]])

    def calculate_g(self):
        self.g = np.zeros(((self.N + 1) * self.shape + self.N * self.shape_u, 1))
        for i in range(self.N + 1):
            if i == 0:
                node = np.dot(-self.qd, self.xk)
            else:
                x_ref = self.x_ref[self.shape * (i - 1): self.shape * i, :]
                node = np.dot(-self.qd, x_ref)
            self.g[self.shape * i:self.shape * i + self.shape, :] = node

    def calculate_Acon(self, lat: bool = True):
        # 构造状态等式约束
        Im = np.eye(self.shape) * 1
        In = np.eye(self.shape_u) * 1
        Aeq = np.kron(np.eye(self.N + 1), - Im)
        if not lat:
            self.calculate_a_lon()
        for i in range(self.N):
            if lat:
                self.calculate_a(self.v_ref[i], self.dt)
            Aeq[self.shape * (i + 1):self.shape * (i + 2), self.shape * i:self.shape * (i + 1)] = self.ad

        # 构造输入约束
        Aeu = np.zeros(((self.N + 1) * self.shape, self.N * self.shape_u))
        for i in range(self.N):
            Aeu[self.shape * (i + 1):self.shape * (i + 2), self.shape_u * i:self.shape_u * (i + 1)] = self.bd

        # 构造状态不等式约束
        Aix = np.kron(np.eye(self.N + 1), Im)
        Aiu = np.kron(np.eye(self.N), In)
        # print(Aeq.shape)
        # print(Aeu.shape)
        # print(Aix.shape)
        # print(Aiu.shape)
        # print(np.zeros(((self.N + 1) * self.shape, self.N * self.shape_u)).shape)
        # print(np.zeros((self.N * self.shape_u, (self.N + 1) * self.shape)).shape)

        self.Acon = np.block([[Aeq, Aeu],
                             [Aix, np.zeros(((self.N + 1) * self.shape, self.N * self.shape_u))],
                             [np.zeros((self.N * self.shape_u, (self.N + 1) * self.shape)), Aiu]])

    def calculate_min_max(self, min_u: float, max_u: float, state_min: list = None, state_max: list = None):
        self.min_u = np.zeros(((self.N + 1) * self.shape * 2 + self.N * self.shape_u, 1))
        self.max_u = np.zeros(((self.N + 1) * self.shape * 2 + self.N * self.shape_u, 1))
        self.min_u[0:self.shape, :] = - self.xk
        self.max_u[0:self.shape, :] = - self.xk
        if state_min is not None and state_max is not None:
            x_min = np.array([[state_min[0]], [state_min[1]], [state_min[2]]])
            x_max = np.array([[state_max[0]], [state_max[1]], [state_max[2]]])
        else:
            x_min = np.array([[-10], [-30], [-30]])
            x_max = np.array([[10], [30], [30]])
        begin_index = (self.N + 1) * self.shape
        for i in range(self.N + 1):
            self.min_u[begin_index + self.shape * i:begin_index + self.shape * (i + 1), :] = x_min
            self.max_u[begin_index + self.shape * i:begin_index + self.shape * (i + 1), :] = x_max
            self.min_u[begin_index * 2 + self.shape_u * i:begin_index * 2 + self.shape_u * (i + 1), :] = min_u
            self.max_u[begin_index * 2 + self.shape_u * i:begin_index * 2 + self.shape_u * (i + 1), :] = max_u

    def test_two(self):
        self.calculate_qn([200, 10, 0])
        self.calculate_rn([20])
        self.calculate_H()
        self.calculate_g()
        self.calculate_Acon()
        self.calculate_min_max(- 30, 30)
        # self.init_qp()
        return self.update_qp()

    def init_qp(self):
        self.qp.setup(P=sparse.csc_matrix(self.H), q=self.g, A=sparse.csc_matrix(self.Acon),
                      l=self.min_u, u=self.max_u, verbose=False)
        self.qp_setup_done = True

    def update_qp(self, lat: bool = False):
        if self.qp_setup_done:
            if lat:
                H_scr = sparse.csr_matrix(self.H)
                A_scr = sparse.csr_matrix(self.Acon)
                PX = H_scr.data
                PX_idx = H_scr.indices
                AX = A_scr.data
                AX_idx = A_scr.indices
                self.qp.update(q=self.g, l=self.min_u, u=self.max_u, Px=PX, Ax=AX, Ax_idx=AX_idx, Px_idx=PX_idx)
            else:
                # try:
                H_csr = sparse.csr_matrix(self.H)
                Px = H_csr.data
                Px_idx = H_csr.indices
                self.qp.update(q=self.g, l=self.min_u, u=self.max_u, Px=Px)
                # self.qp.update(q=self.g, l=self.min_u, u=self.max_u, Px_idx=Px_idx)
            # except:
            #     self.init_qp()
        else:
            self.init_qp()
        self.res = self.qp.solve()
        if self.res.info.status != 'solved':
            logger.error("QP not solved")
            self.qp_failed_time += 1
            if self.qp_failed_time > self.N:
                return None
            if self.predicted_control is not None and any(self.predicted_control[self.qp_failed_time:self.qp_failed_time + 1]):
                return self.predicted_control[self.qp_failed_time:self.qp_failed_time + 1][0]
            return None
        res = self.res.x.reshape(((self.N + 1) * self.shape + self.N * self.shape_u, self.shape_u))
        self.predicted_state = res[self.shape: (self.N + 1) * self.shape]
        self.predicted_control = res[(self.N + 1) * self.shape:]
        res_k = res[(self.N + 1) * self.shape][0]
        self.qp_failed_time = 0
        return res_k

    def solove_qp(self):
        self.res = self.qp.solve()
        if self.res.info.status != 'solved':
            logger.error("QP not solved")
            self.qp_failed_time += 1
            return None
        res = self.res.x.reshape(((self.N + 1) * self.shape + self.N * self.shape_u, self.shape_u))
        self.predicted_state = res[self.shape: (self.N + 1) * self.shape]
        self.predicted_control = res[(self.N + 1) * self.shape:]
        res_k = res[(self.N + 1) * self.shape][0]
        return res_k

    def calculate_x_ref(self, x_ref) -> bool:
        if x_ref is None:
            return False
        self.x_ref = np.zeros((self.N * self.shape, 1))
        for i in range(self.N):
            if i == 0:
                self.x_ref[self.shape * i:self.shape * i + 1, :] = float(x_ref['y'][i]) - float(self.xk[0][0])
                delta_h = float(x_ref['h'][i]) - float(self.xk[1][0])
                if delta_h < -np.pi:
                    delta_h += 2 * np.pi
                elif delta_h > np.pi:
                    delta_h -= 2 * np.pi
                self.x_ref[self.shape * i + 1:self.shape * i + 2, :] = - delta_h
            else:
                self.x_ref[self.shape * i:self.shape * i + 1, :] = float(x_ref['y'][i]) - float(x_ref['y'][i - 1])
                delta_h = float(x_ref['h'][i]) - float(x_ref['h'][i - 1])
                if delta_h < -np.pi:
                    delta_h += 2 * np.pi
                elif delta_h > np.pi:
                    delta_h -= 2 * np.pi
                self.x_ref[self.shape * i + 1:self.shape * i + 2, :] = - delta_h
            self.x_ref[self.shape * i + 2:self.shape * i + 3, :] = 0

    def calculate_x_ref_lon(self, x_ref) -> bool:
        if x_ref is None:
            return False
        self.x_ref = np.zeros((self.N * self.shape, 1))
        for i in range(self.N):
            self.x_ref[self.shape * i:self.shape * i + 1, :] = float(x_ref['s'][i])
            self.x_ref[self.shape * i + 1:self.shape * i + 2, :] = float(x_ref['v'][i])
            self.x_ref[self.shape * i + 2:self.shape * i + 3, :] = 0

    def calculate_H_g(self) -> bool:
        if self.Bd is None:
            return False
        self.H = np.dot(np.dot(self.Bd.T, self.Qd), self.Bd) + self.Rd
        self.g = np.dot(np.dot(np.dot(self.xk.T, self.Ad.T) - self.x_ref.T, self.Qd), self.Bd)
        self.g = np.array(self.g).flatten()
        # self.g = np.dot(np.dot(self.xk.T, np.dot(self.Ad.T, self.Qd)) - self.x_ref.T, self.Bd)

    def calculate_limit_u(self, min_u: float, max_u: float):
        self.min_u = np.zeros((self.N, self.uk.shape[1]))
        self.max_u = np.zeros((self.N, self.uk.shape[1]))
        for i in range(self.N):
            for j in range(self.uk.shape[1]):
                self.min_u[i][j] = min_u
                self.max_u[i][j] = max_u
        self.tmp = np.eye(self.N * self.uk.shape[1]) * 1

    def calculate_res_x(self, res):
        x_res = np.dot(self.Ad, self.xk) + np.dot(self.Bd, res)
        return x_res


    def test_qp(self):
        self.qp.setup(P=sparse.csc_matrix(self.H), q=self.g, A=sparse.csc_matrix(self.tmp),
                      l=self.min_u, u=self.max_u, verbose=False)
        res = self.qp.solve()
        res = res.x.reshape((self.N, self.uk.shape[1]))
        res = res * 180 / np.pi
        return res
